from copy import deepcopy
import algo_utils as utils
import oxc_rest_api as oxc_api
import time
from _typeshed import Self

class networkController():
    TOTAL_CONTROL_RETRY = 5
    RETRY_INTERVAL_MS = 100
    def __init__(self, net_conf):
        self.clos_n = -1
        self.clos_m = -1
        #oxc conf
        self.oxc_map = {}  # {id: {"addr": ip, "port_id_to_seq":{port_id: port_seq}}}
        self.oxc_port_num = -1
        self.dim_lens = []
        #switch conf
        self.leaf_switch_map = {}  # {id: {"addr": ip_or_name, "port_id_to_seq":{port_id: port_seq}, "gpu_addr":{gpu_id: ip}}}
        self.spine_switch_map = {} # {id: {"addr": ip_or_name, "port_id_to_seq":{port_id: port_seq}}}

        self._parse_conf(net_conf)

    def _parse_conf(self, net_conf):
        self.clos_n = net_conf["clos_n"]
        self.clos_m = net_conf["clos_m"]

        #oxc conf
        self.oxc_port_num = net_conf["oxc_port"]
        self.dim_lens.append(net_conf["oxc_dim1_limit"])
        self.dim_lens.append(net_conf["oxc_dim2_limit"])
        self.dim_lens.append(net_conf["oxc_dim3_limit"])
        for info in net_conf["oxc_info"]:
            self.oxc_map[info["id"]] = {}
            self.oxc_map[info["id"]]["addr"] = info["addr"]
            self.oxc_map[info["id"]]["port_id_to_seq"] = {}
            start_seq = info["port_start"]
            end_seq = info["port_end"]
            start_index = utils.oxc_port_seq_to_index(start_seq, self.dim_lens)
            end_index = utils.oxc_port_seq_to_index(end_seq, self.dim_lens)
            for port_id, seq_id in enumerate(range(start_index, end_index+1)):
                self.oxc_map[info["id"]]["port_id_to_seq"][port_id]  = utils.oxc_port_index_to_seq(seq_id, self.dim_lens)

        #switch conf
        # leaf
        for info in net_conf["leaf_switch_info"]:
            self.leaf_switch_map[info["id"]] = {}
            self.leaf_switch_map[info["id"]]["addr"] = info["addr"]
            self.leaf_switch_map[info["id"]]["port_id_to_seq"] = {}
            self.leaf_switch_map[info["id"]]["gpu_addr"] = {}
            start_seq = info["port_start"]
            end_seq = info["port_end"]
            start_index = utils.switch_port_seq_to_index(start_seq)
            end_index = utils.switch_port_seq_to_index(end_seq)
            for port_id, seq_id in enumerate(range(start_index, end_index + 1)):
                self.leaf_switch_map[info["id"]]["port_id_to_seq"][port_id] = utils.switch_port_index_to_seq(seq_id)
            for gpu_id, gpu_addr in enumerate(info["gpu_addr"]):
                self.leaf_switch_map[info["id"]]["gpu_addr"][gpu_id] = gpu_addr

        # spine
        for info in net_conf["spine_switch_info"]:
            self.spine_switch_map[info["id"]] = {}
            self.spine_switch_map[info["id"]]["addr"] = info["addr"]
            self.spine_switch_map[info["id"]]["port_id_to_seq"] = {}
            start_seq = info["port_start"]
            end_seq = info["port_end"]
            start_index = utils.switch_port_seq_to_index(start_seq)
            end_index = utils.switch_port_seq_to_index(end_seq)
            for port_id, seq_id in enumerate(range(start_index, end_index + 1)):
                self.spine_switch_map[info["id"]]["port_id_to_seq"][port_id] = utils.switch_port_index_to_seq(seq_id)

    def switch_elec(self, leaf_to_spine_linkage, gpu_indexes):# [spine_id, spine_port_id, leaf_id], indexes of the allocated gpus (global index)
        if len(leaf_to_spine_linkage) == 0: return True

        routing_material_leaf = {} # {leaf_switch_id: {"gpu_addr":[], "leaf_up_port_seq":[], "spine_port_seq":[(spine_switch_id, spine_port_seq)]}}
        routing_material_spine = {} # {spine_switch_id: {"spine_port_seq":[], "leaf_port_seq":[(leaf_switch_id, leaf_port_seq)]}}
        # note that the mapping from gpu_addr to leaf_up_port haven't been decided yet
        # we leave it for the routing algorithm to decide


        # *********************************************
        # mapping link entries to routing rules
        # *********************************************
        for linkage in leaf_to_spine_linkage:
            spine_id = linkage[0]
            spine_port_id = linkage[1] # decide oxc index -> decide leaf up port
            leaf_id = linkage[2]
            oxc_id = int(spine_port_id / 2)
            up_port_id =int(len(self.leaf_switch_map[leaf_id]["port_id_to_seq"].keys()) /2) + oxc_id

            #fill leaf routing material
            if leaf_id not in routing_material_leaf.keys():
                routing_material_leaf[leaf_id] =  {"gpu_addr":[], "leaf_up_port_seq":[], "spine_port_seq":[]}
            routing_material_leaf[leaf_id]["leaf_up_port_seq"].append(self.leaf_switch_map[leaf_id]["port_id_to_seq"][up_port_id])
            routing_material_leaf[leaf_id]["spine_port_seq"].append((spine_id, self.spine_switch_map[spine_id]["port_id_to_seq"][spine_port_id]))

            #fill spine routing material
            if spine_id not in routing_material_spine.keys():
                routing_material_spine[spine_id] = {"spine_port_seq":[], "leaf_port_seq":[]}
            routing_material_spine[spine_id]["spine_port_seq"].append(self.spine_switch_map[spine_id]["port_id_to_seq"][spine_port_id])
            routing_material_spine[spine_id]["leaf_port_seq"].append((leaf_id, self.leaf_switch_map[leaf_id]["port_id_to_seq"][up_port_id]))
        #fill gpu indexes
        for g_id in gpu_indexes:
            # translate global index to leaf_id + local
            leaf_id = int(g_id/self.clos_n)
            local_id = g_id % self.clos_n
            assert leaf_id in routing_material_leaf.keys()
            routing_material_leaf[leaf_id]["gpu_addr"].append(self.leaf_switch_map[leaf_id]["port_id_to_seq"][local_id])


        # tricky part: need to deal with cases where multiple n/m exists in the same leaf/spine
        # *************************************
        #  Routing 1: source-based up routing in leaf
        # *************************************
        pass

        # *************************************
        #  Routing 2: src-dest-based routing in spine
        # *************************************
        pass

        # *************************************
        #  Routing 3: dest-based down routing in leaf
        # *************************************
        pass

    def switch_oxc(self, leaf_to_spine_linkage):# [spine_id, spine_port_id, leaf_id]
        if len(leaf_to_spine_linkage) == 0: return True

        oxc_ports_to_connect = [] # ipaddr, port_seq1, port_seq2
        for linkage in leaf_to_spine_linkage:
            spine_id = linkage[0]
            spine_port_id  = linkage[1]
            leaf_id = linkage[2]
            # translate to oxc connection
            up_port_id = int(self.oxc_port_num/2) + spine_id
            down_port_id = leaf_id
            oxc_id = int(spine_port_id/2)
            oxc_ports_to_connect.append([self.oxc_map[oxc_id]["addr"], self.oxc_map[oxc_id]["port_id_to_seq"][down_port_id], self.oxc_map[oxc_id]["port_id_to_seq"][up_port_id],])

        retried_times =0
        for entry in oxc_ports_to_connect:
            res = eval(oxc_api.add_connection(entry[1], entry[2], entry[0]))
            while res["status"] != 200 and retried_times < self.TOTAL_CONTROL_RETRY:
                time.sleep(self.RETRY_INTERVAL_MS / 1000.0)
                res = eval(oxc_api.add_connection(entry[1], entry[2], entry[0]))
                retried_times += 1
            if res["status"] != 200:
                return False

        return True


# defined

